"""
预览数据集

- 按下n键预览下一张图片
- 否则退出
"""

import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

label = pd.read_csv("./data/train_solution_bounding_boxes.csv")
print(label.shape)


# read images list
data_path = './data/training_images/'
images = label.image.unique()
np.random.shuffle(images)

continue_key = 0
while continue_key != -1:
    # multiple car in one image
    image_name = images[continue_key]
    samples = label[label.image == image_name]
    image = cv2.imread(os.path.join(data_path, image_name))
    print(image.shape)
    for i in samples.index:
        label_i = samples.loc[i]
        xmin, ymin, xmax, ymax = label_i['xmin'], label_i['ymin'], label_i['xmax'], label_i['ymax']
        # cast int
        xmin = int(np.round(xmin))
        ymin = int(np.round(ymin))
        xmax = int(np.round(xmax))
        ymax = int(np.round(ymax))

        points = [
            (xmin, ymin),
            (xmax, ymin),
            (xmax, ymax),
            (xmin, ymax)
        ]
        for i in range(4):
            cv2.line(image, points[i], points[(i + 1) % 4], (255, 0, 0), 3)
    cv2.imshow(image_name, image)
    key = cv2.waitKey(0)
    cv2.destroyAllWindows()
    print(key)
    # plt.figure("Image")
    # plt.imshow(image)
    # plt.title(image_name)
    # plt.show()
    # plt.clf()
    continue_key = -1 if key != ord('n') else continue_key + 1
    # continue_key = -1

# cv2.destroyAllWindows()
